// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/gtest_utils.h>
#include <pollux/common/testutil/optional_empty.h>
#include <pollux/exec/plan_node_stats.h>
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/testing/exec/util/operator_test_base.h>
#include <pollux/plan/plan_builder.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::exec::test;

class UnnestTest : public OperatorTestBase,
                   public testing::WithParamInterface<vector_size_t> {
  void SetUp() override {
    OperatorTestBase::SetUp();
  }

  void TearDown() override {
    OperatorTestBase::TearDown();
  }

 protected:
  CursorParameters makeCursorParameters(
      const core::PlanNodePtr& planNode) const {
    CursorParameters params;
    params.planNode = planNode;
    params.queryConfigs[core::QueryConfig::kPreferredOutputBatchRows] =
        std::to_string(batchSize_);
    return params;
  }

  const vector_size_t batchSize_{GetParam()};
};

TEST_P(UnnestTest, basicArray) {
  auto vector = make_row_vector({
      make_flat_vector<int64_t>(100, [](auto row) { return row; }),
      make_array_vector<int32_t>(
          100,
          [](auto row) { return row % 5 + 1; },
          [](auto row, auto index) { return index * (row % 3); },
          null_every(7)),
  });

  createDuckDbTable({vector});

  // TODO Add tests with empty arrays. This requires better support in DuckDB.

  auto op = PlanBuilder().values({vector}).unnest({"c0"}, {"c1"}).planNode();
  auto params = makeCursorParameters(op);
  assertQuery(params, "SELECT c0, UNNEST(c1) FROM tmp WHERE c0 % 7 > 0");
}

TEST_P(UnnestTest, arrayWithOrdinality) {
  auto array = vectorMaker_.array_vector_nullable<int32_t>(
      {{{1, 2, std::nullopt, 4}},
       std::nullopt,
       {{5, 6}},
       common::testutil::optionalEmpty,
       {{{{std::nullopt}}}},
       {{7, 8, 9}}});
  auto vector = make_row_vector(
      {make_nullable_flat_vector<double>({1.1, 2.2, 3.3, 4.4, 5.5, std::nullopt}),
       array});

  auto op = PlanBuilder()
                .values({vector})
                .unnest({"c0"}, {"c1"}, "ordinal")
                .planNode();

  auto expected = make_row_vector(
      {make_nullable_flat_vector<double>(
           {1.1,
            1.1,
            1.1,
            1.1,
            3.3,
            3.3,
            5.5,
            std::nullopt,
            std::nullopt,
            std::nullopt}),
       make_nullable_flat_vector<int32_t>(
           {1, 2, std::nullopt, 4, 5, 6, std::nullopt, 7, 8, 9}),
       make_nullable_flat_vector<int64_t>({1, 2, 3, 4, 1, 2, 1, 1, 2, 3})});
  auto params = makeCursorParameters(op);
  assertQuery(params, expected);

  // Test with array wrapped in dictionary.
  auto reversedIndices = make_indices_in_reverse(6);
  auto vectorInDict = make_row_vector(
      {make_nullable_flat_vector<double>({1.1, 2.2, 3.3, 4.4, 5.5, std::nullopt}),
       wrap_in_dictionary(reversedIndices, 6, array)});
  op = PlanBuilder()
           .values({vectorInDict})
           .unnest({"c0"}, {"c1"}, "ordinal")
           .planNode();

  auto expectedInDict = make_row_vector(
      {make_nullable_flat_vector<double>(
           {1.1,
            1.1,
            1.1,
            2.2,
            4.4,
            4.4,
            std::nullopt,
            std::nullopt,
            std::nullopt,
            std::nullopt}),
       make_nullable_flat_vector<int32_t>(
           {7, 8, 9, std::nullopt, 5, 6, 1, 2, std::nullopt, 4}),
       make_nullable_flat_vector<int64_t>({1, 2, 3, 1, 1, 2, 1, 2, 3, 4})});
  params = makeCursorParameters(op);
  assertQuery(params, expectedInDict);
}

TEST_P(UnnestTest, basicMap) {
  auto vector = make_row_vector(
      {make_flat_vector<int64_t>(100, [](auto row) { return row; }),
       make_map_vector<int64_t, double>(
           100,
           [](auto /* row */) { return 2; },
           [](auto row) { return row % 2; },
           [](auto row) { return row % 2 + 1; })});
  auto op = PlanBuilder().values({vector}).unnest({"c0"}, {"c1"}).planNode();
  // DuckDB doesn't support Unnest from MAP column. Hence,using 2 separate array
  // columns with the keys and values part of the MAP to validate.
  auto duckDbVector = make_row_vector(
      {make_flat_vector<int64_t>(100, [](auto row) { return row; }),
       make_array_vector<int32_t>(
           100,
           [](auto /* row */) { return 2; },
           [](auto /* row */, auto index) { return index; }),
       make_array_vector<int32_t>(
           100,
           [](auto /* row */) { return 2; },
           [](auto /* row */, auto index) { return index + 1; })});
  createDuckDbTable({duckDbVector});
  auto params = makeCursorParameters(op);
  assertQuery(params, "SELECT c0, UNNEST(c1), UNNEST(c2) FROM tmp");
}

TEST_P(UnnestTest, mapWithOrdinality) {
  auto map = make_map_vector<int32_t, double>(
      {{{1, 1.1}, {2, std::nullopt}},
       {{3, 3.3}, {4, 4.4}, {5, 5.5}},
       {{6, std::nullopt}}});
  auto vector =
      make_row_vector({make_nullable_flat_vector<int32_t>({1, 2, 3}), map});

  auto op = PlanBuilder()
                .values({vector})
                .unnest({"c0"}, {"c1"}, "ordinal")
                .planNode();

  auto expected = make_row_vector(
      {make_nullable_flat_vector<int32_t>({1, 1, 2, 2, 2, 3}),
       make_nullable_flat_vector<int32_t>({1, 2, 3, 4, 5, 6}),
       make_nullable_flat_vector<double>(
           {1.1, std::nullopt, 3.3, 4.4, 5.5, std::nullopt}),
       make_nullable_flat_vector<int64_t>({1, 2, 1, 2, 3, 1})});
  auto params = makeCursorParameters(op);
  assertQuery(params, expected);

  // Test with map wrapped in dictionary.
  auto reversedIndices = make_indices_in_reverse(3);
  auto vectorInDict = make_row_vector(
      {make_nullable_flat_vector<int32_t>({1, 2, 3}),
       wrap_in_dictionary(reversedIndices, 3, map)});
  op = PlanBuilder()
           .values({vectorInDict})
           .unnest({"c0"}, {"c1"}, "ordinal")
           .planNode();

  auto expectedInDict = make_row_vector(
      {make_nullable_flat_vector<int32_t>({1, 2, 2, 2, 3, 3}),
       make_nullable_flat_vector<int32_t>({6, 3, 4, 5, 1, 2}),
       make_nullable_flat_vector<double>(
           {std::nullopt, 3.3, 4.4, 5.5, 1.1, std::nullopt}),
       make_nullable_flat_vector<int64_t>({1, 1, 2, 3, 1, 2})});
  params = makeCursorParameters(op);
  assertQuery(params, expectedInDict);
}

TEST_P(UnnestTest, multipleColumns) {
  std::vector<vector_size_t> offsets(100, 0);
  for (int i = 1; i < 100; ++i) {
    offsets[i] = offsets[i - 1] + i % 11 + 1;
  }

  auto vector = make_row_vector(
      {make_flat_vector<int64_t>(100, [](auto row) { return row; }),
       vectorMaker_.map_vector<int64_t, double>(
           100,
           [](auto row) { return row % 7 + 1; },
           [](auto row, auto index) { return index * (row % 4); },
           [](auto row, auto index) { return index * (row % 4) + 1; },
           null_every(6)),
       make_array_vector<int32_t>(
           100,
           [](auto row) { return row % 5 + 1; },
           [](auto row, auto index) { return index * (row % 3); },
           null_every(7)),
       make_array_vector(offsets, make_constant<int32_t>(7, 700))});

  auto op = PlanBuilder()
                .values({vector})
                .unnest({"c0"}, {"c1", "c2", "c3"})
                .planNode();

  // DuckDB doesn't support Unnest from MAP column. Hence,using 2 separate array
  // columns with the keys and values part of the MAP to validate.
  auto duckDbVector = make_row_vector(
      {make_flat_vector<int64_t>(100, [](auto row) { return row; }),
       make_array_vector<int64_t>(
           100,
           [](auto row) { return row % 7 + 1; },
           [](auto row, auto index) { return index * (row % 4); },
           null_every(6)),
       make_array_vector<double>(
           100,
           [](auto row) { return row % 7 + 1; },
           [](auto row, auto index) { return index * (row % 4) + 1; },
           null_every(6)),
       make_array_vector<int32_t>(
           100,
           [](auto row) { return row % 5 + 1; },
           [](auto row, auto index) { return index * (row % 3); },
           null_every(7)),
       make_array_vector(offsets, make_constant<int32_t>(7, 700))});
  createDuckDbTable({duckDbVector});
  auto params = makeCursorParameters(op);
  assertQuery(
      params,
      "SELECT c0, UNNEST(c1), UNNEST(c2), UNNEST(c3), UNNEST(c4) FROM tmp");
}

TEST_P(UnnestTest, multipleColumnsWithOrdinality) {
  std::vector<vector_size_t> offsets(100, 0);
  for (int i = 1; i < 100; ++i) {
    offsets[i] = offsets[i - 1] + i % 11 + 1;
  }

  auto vector = make_row_vector(
      {make_flat_vector<int64_t>(100, [](auto row) { return row; }),
       vectorMaker_.map_vector<int64_t, double>(
           100,
           [](auto row) { return row % 7 + 1; },
           [](auto row, auto index) { return index * (row % 4); },
           [](auto row, auto index) { return index * (row % 4) + 1; },
           null_every(6)),
       make_array_vector<int32_t>(
           100,
           [](auto row) { return row % 5 + 1; },
           [](auto row, auto index) { return index * (row % 3); },
           null_every(7)),
       make_array_vector(offsets, make_constant<int32_t>(7, 700))});

  auto op = PlanBuilder()
                .values({vector})
                .unnest({"c0"}, {"c1", "c2", "c3"}, "ordinal")
                .planNode();

  // DuckDB doesn't support Unnest from MAP column. Hence,using 2 separate array
  // columns with the keys and values part of the MAP to validate.
  auto ordinalitySize = [&](auto row) {
    if (row % 42 == 0) {
      return offsets[row + 1] - offsets[row];
    } else if (row % 7 == 0) {
      return std::max(row % 7 + 1, offsets[row + 1] - offsets[row]);
    } else if (row % 6 == 0) {
      return std::max(row % 5 + 1, offsets[row + 1] - offsets[row]);
    } else {
      return std::max(
          std::max(row % 5, row % 7) + 1,
          (row == 99 ? 700 : offsets[row + 1]) - offsets[row]);
    }
  };

  auto duckDbVector = make_row_vector(
      {make_flat_vector<int64_t>(100, [](auto row) { return row; }),
       make_array_vector<int64_t>(
           100,
           [](auto row) { return row % 7 + 1; },
           [](auto row, auto index) { return index * (row % 4); },
           null_every(6)),
       make_array_vector<double>(
           100,
           [](auto row) { return row % 7 + 1; },
           [](auto row, auto index) { return index * (row % 4) + 1; },
           null_every(6)),
       make_array_vector<int32_t>(
           100,
           [](auto row) { return row % 5 + 1; },
           [](auto row, auto index) { return index * (row % 3); },
           null_every(7)),
       make_array_vector(offsets, make_constant<int32_t>(7, 700)),
       make_array_vector<int64_t>(
           100, ordinalitySize, [](auto /* row */, auto index) {
             return index + 1;
           })});
  createDuckDbTable({duckDbVector});
  auto params = makeCursorParameters(op);
  assertQuery(
      params,
      "SELECT c0, UNNEST(c1), UNNEST(c2), UNNEST(c3), UNNEST(c4), UNNEST(c5) FROM tmp");

  // Test with empty arrays and maps.
  vector = make_row_vector(
      {make_nullable_flat_vector<double>({1.1, 2.2, std::nullopt, 4.4, 5.5}),
       vectorMaker_.array_vector_nullable<int32_t>(
           {{{1, 2, std::nullopt, 4}}, std::nullopt, {{5, 6}}, {}, {{7}}}),
       make_map_vector<int32_t, double>(
           {{{1, 1.1}, {2, std::nullopt}},
            {{3, 3.3}, {4, 4.4}, {5, 5.5}},
            {{6, std::nullopt}},
            {},
            {}})});

  op = PlanBuilder()
           .values({vector})
           .unnest({"c0"}, {"c1", "c2"}, "ordinal")
           .planNode();

  auto expected = make_row_vector(
      {make_nullable_flat_vector<double>(
           {1.1,
            1.1,
            1.1,
            1.1,
            2.2,
            2.2,
            2.2,
            std::nullopt,
            std::nullopt,
            5.5}),
       make_nullable_flat_vector<int32_t>(
           {1,
            2,
            std::nullopt,
            4,
            std::nullopt,
            std::nullopt,
            std::nullopt,
            5,
            6,
            7}),
       make_nullable_flat_vector<int32_t>(
           {1,
            2,
            std::nullopt,
            std::nullopt,
            3,
            4,
            5,
            6,
            std::nullopt,
            std::nullopt}),
       make_nullable_flat_vector<double>(
           {1.1,
            std::nullopt,
            std::nullopt,
            std::nullopt,
            3.3,
            4.4,
            5.5,
            std::nullopt,
            std::nullopt,
            std::nullopt}),
       make_nullable_flat_vector<int64_t>({1, 2, 3, 4, 1, 2, 3, 1, 2, 1})});
  params = makeCursorParameters(op);
  assertQuery(params, expected);
}

TEST_P(UnnestTest, allEmptyOrNullArrays) {
  auto vector = make_row_vector(
      {make_flat_vector<int64_t>(100, [](auto row) { return row; }),
       make_array_vector<int32_t>(
           100,
           [](auto /* row */) { return 0; },
           [](auto /* row */, auto index) { return index; },
           null_every(5)),
       make_array_vector<int32_t>(
           100,
           [](auto /* row */) { return 0; },
           [](auto /* row */, auto index) { return index; },
           null_every(7))});

  auto op =
      PlanBuilder().values({vector}).unnest({"c0"}, {"c1", "c2"}).planNode();
  auto params = makeCursorParameters(op);
  assertQueryReturnsEmptyResult(params);

  op = PlanBuilder()
           .values({vector})
           .unnest({"c0"}, {"c1", "c2"}, "ordinal")
           .planNode();
  assertQueryReturnsEmptyResult(params);
}

TEST_P(UnnestTest, allEmptyOrNullMaps) {
  auto vector = make_row_vector(
      {make_flat_vector<int64_t>(100, [](auto row) { return row; }),
       make_map_vector<int64_t, double>(
           100,
           [](auto /* row */) { return 0; },
           [](auto /* row */) { return 0; },
           [](auto /* row */) { return 0; },
           null_every(5)),
       make_map_vector<int64_t, double>(
           100,
           [](auto /* row */) { return 0; },
           [](auto /* row */) { return 0; },
           [](auto /* row */) { return 0; },
           null_every(7))});

  auto op =
      PlanBuilder().values({vector}).unnest({"c0"}, {"c1", "c2"}).planNode();
  auto params = makeCursorParameters(op);
  assertQueryReturnsEmptyResult(params);

  op = PlanBuilder()
           .values({vector})
           .unnest({"c0"}, {"c1", "c2"}, "ordinal")
           .planNode();
  params = makeCursorParameters(op);
  assertQueryReturnsEmptyResult(params);
}

TEST_P(UnnestTest, batchSize) {
  auto data = make_row_vector({
      make_flat_vector<int64_t>(1'000, [](auto row) { return row; }),
  });

  // Unnest 1K rows into 3K rows.
  core::PlanNodeId unnestId;
  auto plan = PlanBuilder()
                  .values({data})
                  .project({"sequence(1, 3) as s"})
                  .unnest({}, {"s"})
                  .capturePlanNodeId(unnestId)
                  .planNode();

  auto expected = make_row_vector({
      make_flat_vector<int64_t>(1'000 * 3, [](auto row) { return 1 + row % 3; }),
  });

  auto task = AssertQueryBuilder(plan)
                  .config(
                      core::QueryConfig::kPreferredOutputBatchRows,
                      std::to_string(batchSize_))
                  .assertResults({expected});
  auto stats = exec::toPlanStats(task->taskStats());

  ASSERT_EQ(3'000, stats.at(unnestId).outputRows);
  int32_t expectedNumVectors = 3'000 / batchSize_;
  if (3'000 % batchSize_ != 0) {
    expectedNumVectors++;
  }
  ASSERT_EQ(expectedNumVectors, stats.at(unnestId).outputVectors);
}

POLLUX_INSTANTIATE_TEST_SUITE_P(
    UnnestTest,
    UnnestTest,
    testing::ValuesIn(/*batchSize*/ {2, 17, 33, 1024}));
